# -*- coding: utf-8 -*-
"""
Created on Tue Oct 26 18:29:17 2021

@author: think
"""

from pyDes import des, PAD_PKCS5
from hashlib import md5
import rsa



class digitalEnvelope:
    def __init__(self):
        pass
    
    
    #数字信封
    def getEnvelope(self, message, privkeyS, pubkeyR, key):
        #输入：message：一般字符串，privkeyS：rsa库的private_key类，
        #pubkeyR：接收方的公钥，rsa库的private_key类，有点拧巴但必须这么写
        #key：一般8位字符串，输出：utf-8字符串
        
        Sign = self.digitalSignature(message, privkeyS)
        
        midMessage = message.encode('utf-8') + Sign
        crypto_msg = self.desEncoding(midMessage, key)
        crypto_key = self.rsaEncryption(key.encode("utf-8"), pubkeyR)
        
        crypto = crypto_msg + crypto_key
        return crypto
        
      
    def removeEnvelope(self, crypto, pubkeyS, privkeyR):
        #获得公钥key
        crypto_key = crypto[-64:]#后64位为公钥RSA加密后的结果
        crypto_msg = crypto[:-64]
        key = self.rsaDecryption(crypto_key, privkeyR)
        
        #获得消息和数字签名
        midMessage = self.desDecoding(crypto_msg, key)
        Sign = midMessage[-64:]#数字签名也一定是64位
        message = midMessage[:-64]
        
        #有效性验证
        #md5不该打包到digitalSignature里来着，失策了，但不想改了
        m = md5()
        m.update(message)
        checkHash = m.digest()
        
        msgHash = self.rsaDecryption(Sign, pubkeyS)
        if msgHash == checkHash:
            print("The message has not been changed.")
        else:
            print("The message has been changed.")
        
        return message
    
    
    #消息摘要
    def digitalSignature(self, message, privkeyS):#利用Sender的私钥加密
        #输入：message：一般字符串，privkeyS：int; 输出：64位utf-8字符串
        
        #先获得摘要
        m = md5()
        m.update(message.encode('utf-8'))
        Hash = m.digest()
        #这里获得utf-8字符串
        #然后用Sender的私钥加密该摘要
        Sign = self.rsaEncryption(Hash, privkeyS)
        
        return Sign
    
    
    #获得RSA密钥
    def generateRSAKeys(self, num = 512):
        return rsa.newkeys(num)
    
    #RSA加密解密
    #在这里，ke和kd并不一定是公钥或私钥
    #默认采用512位RSA加密，即128位16进制数
    def rsaEncryption(self, message, ke):
        #输入：任意utf-8字符串，输出：64位utf-8字符串
        C = rsa.encrypt(message, ke)
        
        return C
    
    def rsaDecryption(self, message, kd):
        #输入：64位utf-8字符串，输出：任意utf-8字符串
        M = rsa.decrypt(message, kd)
        
        return M
    

    #Des加密解密    
    def desEncoding(self, message, key):
        #输入：message：任意utf-8字符串，key：8位任意字符串; 输出：8的整数倍位数utf-8字符串
        d = des(key, padmode = PAD_PKCS5)
        C = d.encrypt(message, padmode = PAD_PKCS5)
        
        return C
    
    def desDecoding(self, C, key):
        #输入：C：utf-8字符串，key：8位任意字符串; 输出：任意utf-8字符串
        d = des(key, padmode = PAD_PKCS5)
        M = d.decrypt(C, padmode = PAD_PKCS5)
        
        return M
